#!/usr/bin/env python3
"""
使用 Whisper 将音频转为文本的工具 - 命令行版（中文界面）
工具：pip install openai-whisper tqdm
依赖：FFMPEG，如需CUDA加速，需PyTorch

功能亮点：
- 自动使用 GPU（如可用），显著加速转录
- 支持单个或批量音频文件转录
- 输出格式：TXT（文本） + SRT（字幕，可选）
- 自动或手动指定语言（如 'zh', 'en', 'ja' 等）
- 实时进度条 + 详细中文日志
- 错误隔离：单个文件失败不影响整体流程
- 支持自定义模型目录（--model-root）

使用示例：
  python speak_to_txt.py 录音.wav
  python speak_to_txt.py --model small --language zh *.mp3
  python speak_to_txt.py 会议1.wav 会议2.mp3 --srt --model-root ./models
"""

import whisper
import os
import time
import argparse
import logging
from pathlib import Path
from typing import List, Optional, Dict, Any
from tqdm import tqdm
import torch


# 配置中文日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger(__name__)


class AudioTranscriber:
    def __init__(self, model_size: str = "turbo", model_root: Optional[str] = None):
        """初始化转录器，自动检测并使用 GPU（如可用）"""
        # 检测设备
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"检测到设备: {'GPU (CUDA)' if self.device == 'cuda' else 'CPU'}")

        # 如果指定了 model_root，确保目录存在（避免下载时出错）
        if model_root is not None:
            os.makedirs(model_root, exist_ok=True)

        logger.info(f"正在加载 Whisper '{model_size}' 模型（设备: {self.device.upper()}）...")
        self.model = whisper.load_model(
            model_size,
            device=self.device,
            download_root=model_root  # None 或 有效目录路径
        )
        logger.info("✅ 模型加载成功！")

    def transcribe_file(self, audio_path: str, language: Optional[str] = None) -> Dict[str, Any]:
        """
        转录单个音频文件

        参数:
            audio_path: 音频文件路径
            language: 语言代码（如 'zh', 'en', 'ja'），None 表示自动检测
        """
        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"音频文件不存在: {audio_path}")

        audio_name = Path(audio_path).name
        logger.info(f"开始转录: {audio_name}")

        start_time = time.time()
        options = {"language": language} if language else {}
        try:
            result = self.model.transcribe(audio_path, **options)
        except Exception as e:
            logger.error(f"转录失败 ({audio_name}): {e}")
            raise

        processing_time = time.time() - start_time
        detected_lang = result.get('language', '未知')
        logger.info(f"✅ 完成! 耗时: {processing_time:.1f} 秒 | 检测语言: {detected_lang}")

        return {
            'text': result['text'].strip(),
            'language': detected_lang,
            'segments': result.get('segments', []),
            'processing_time': processing_time,
            'audio_path': audio_path
        }

    def save_transcript_txt(self, result: Dict[str, Any], output_path: Path) -> None:
        """保存转录文本为 .txt 文件"""
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write("=== 音频转录结果 ===\n")
            f.write(f"音频文件: {Path(result['audio_path']).name}\n")
            f.write(f"识别语言: {result['language']}\n")
            f.write(f"处理耗时: {result['processing_time']:.1f} 秒\n")
            f.write("=" * 40 + "\n\n")
            f.write(result['text'])
        logger.info(f"✅ 文本已保存: {output_path}")

    def save_subtitles_srt(self, result: Dict[str, Any], output_path: Path) -> None:
        """保存带时间戳的字幕为 .srt 文件"""
        def format_timestamp(seconds: float) -> str:
            """将秒数转换为 SRT 时间戳格式 (HH:MM:SS,mmm)"""
            hrs = int(seconds // 3600)
            mins = int((seconds % 3600) // 60)
            secs = seconds % 60
            ms = int((secs - int(secs)) * 1000)
            secs = int(secs)
            return f"{hrs:02d}:{mins:02d}:{secs:02d},{ms:03d}"

        with open(output_path, 'w', encoding='utf-8') as f:
            for i, segment in enumerate(result['segments'], start=1):
                start = format_timestamp(segment['start'])
                end = format_timestamp(segment['end'])
                text = segment['text'].strip().replace('\n', ' ')
                f.write(f"{i}\n{start} --> {end}\n{text}\n\n")
        logger.info(f"✅ 字幕已保存: {output_path}")


def transcribe_files(
    audio_paths: List[str],
    model_size: str = "turbo",
    language: Optional[str] = None,
    output_srt: bool = False,
    model_root: Optional[str] = None
) -> List[Dict[str, Any]]:
    """批量转录音频文件，带进度条"""
    transcriber = AudioTranscriber(model_size=model_size, model_root=model_root)
    results = []

    for audio_path in tqdm(audio_paths, desc="正在转录", unit="文件"):
        try:
            result = transcriber.transcribe_file(audio_path, language=language)
            results.append(result)

            # 保存 TXT
            audio_stem = Path(audio_path).stem
            txt_path = Path(f"{audio_stem}_转录文本.txt")
            transcriber.save_transcript_txt(result, txt_path)

            # 保存 SRT（如果启用）
            if output_srt:
                srt_path = Path(f"{audio_stem}.srt")
                transcriber.save_subtitles_srt(result, srt_path)

        except Exception as e:
            logger.error(f"跳过文件（因错误）: {Path(audio_path).name} - {e}")
            continue

    return results


def main():
    parser = argparse.ArgumentParser(
        description="使用 OpenAI Whisper 转录音频文件（中文界面）",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        "audio_files",
        nargs="+",
        help="一个或多个音频文件路径（支持通配符，如 *.mp3）"
    )
    parser.add_argument(
        "--model",
        default="turbo",
        choices=["tiny", "base", "small", "medium", "large-v2", "large", "turbo"],
        help="Whisper 模型大小（越大越准，但越慢）\n"
             "tiny/base: 快速测试\n"
             "small/medium: 平衡选择\n"
             "large: 高精度,分为 large-v2和 large-v3，large == large-v3 \n"
             "turbo: 基于large-v3的turbo版，体积比large小（1.58G），turbo == large-v3-trubo,（推荐中文用）"
    )
    parser.add_argument(
        "--language",
        type=str,
        default="zh",
        help="指定语言代码（如：zh=中文, en=英文, ja=日文）\n"
             "不指定则自动检测语言"
    )
    parser.add_argument(
        "--srt",
        action="store_true",
        help="同时生成 .srt 字幕文件（带时间戳）"
    )
    parser.add_argument(
        "--model-root",
        type=str,
        default=None,
        help="指定 Whisper 模型的本地缓存目录（例如：./models）。\n"
             "Whisper 将在此目录中查找或下载 {model}.pt 文件。\n"
             "若未指定，则使用默认缓存路径（~/.cache/whisper）。"
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="显示详细调试日志"
    )

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # 展开通配符（如 *.wav）
    expanded_paths = []
    for pattern in args.audio_files:
        if "*" in pattern or "?" in pattern:
            matches = list(Path().glob(pattern))
        else:
            p = Path(pattern)
            matches = [p] if p.exists() else []
        if not matches:
            logger.warning(f"未找到匹配的文件: {pattern}")
        expanded_paths.extend(str(p.resolve()) for p in matches)

    if not expanded_paths:
        logger.error("❌ 未提供有效的音频文件，请检查路径或通配符。")
        return

    logger.info(f"共找到 {len(expanded_paths)} 个音频文件，开始转录...")

    try:
        results = transcribe_files(
            audio_paths=expanded_paths,
            model_size=args.model,
            language=args.language,
            output_srt=args.srt,
            model_root=args.model_root  # ← 关键：传递 model_root
        )
        logger.info(f"🎉 全部完成！成功转录 {len(results)} 个文件。")
    except KeyboardInterrupt:
        logger.info("⚠️ 用户中断转录过程。")
    except Exception as e:
        logger.error(f"❌ 发生未预期错误: {e}")
        raise


if __name__ == "__main__":
    main()